{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 回溯"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm.script import relax as R\n",
    "from tvm.script import tir as T\n",
    "from tvm import relax as rx\n",
    "from tvm import relay, tir\n",
    "from tvm.relax.analysis import get_var2val\n",
    "import tvm.testing\n",
    "from tvm.relax.dpl import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_backtrack_if_rewriter_returns_no_op():\n",
    "    \"\"\"Rewriter participates in the pattern matching\n",
    "\n",
    "    Sometimes, the pattern-matching syntax is insufficient to check if\n",
    "    a replacement may be performed.  In this case, the `rewriter`\n",
    "    function may perform additional validation.  If this validation\n",
    "    fails, the `rewriter` function can return the original expression,\n",
    "    and no replacement is performed.\n",
    "\n",
    "    In addition, when the `rewriter` returns the original expression,\n",
    "    the pattern match should backtrack to determine if another branch\n",
    "    of the match may have produced a replacement.\n",
    "\n",
    "    This functionality allows pattern replacements to be composed.\n",
    "    \"\"\"\n",
    "\n",
    "    pat_match_no_rewrite = is_op(\"relax.add\")(wildcard(), wildcard())\n",
    "\n",
    "    pat_arg = wildcard()\n",
    "    pat_zeros = is_op(\"relax.zeros\")(wildcard())\n",
    "    pat_add = is_op(\"relax.add\")(pat_arg, pat_zeros)\n",
    "\n",
    "    # OR conditions are checked in the order that they occur.  Because\n",
    "    # `pat_match_no_rewrite` is a superset of `pat_add`, it will\n",
    "    # always match first.\n",
    "    pat = pat_match_no_rewrite | pat_add\n",
    "\n",
    "    def rewriter(expr, matches):\n",
    "        if pat_match_no_rewrite in matches:\n",
    "            # This branch simulates a rewrite whose precondition has\n",
    "            # failed.  If the pattern-matching treats this as a\n",
    "            # successful match with no replacemen required, then no\n",
    "            # rewrite would be performed.  On the other hand, if the\n",
    "            # pattern-matching treats this as an unsuccessful match,\n",
    "            # then it can backtrack and attempt `pat_add` instead.\n",
    "            return expr\n",
    "        elif pat_add in matches:\n",
    "            return matches[pat_arg]\n",
    "        else:\n",
    "            raise RuntimeError(\"Pattern matched, but neither branch matched\")\n",
    "\n",
    "    @R.function(private=True)\n",
    "    def before():\n",
    "        with R.dataflow():\n",
    "            A = R.ones([64, 128], \"int32\")\n",
    "            B = R.zeros([64, 128], \"int32\")\n",
    "            C = R.add(A, B)\n",
    "\n",
    "            R.output(C)\n",
    "        return C\n",
    "\n",
    "    @R.function(private=True)\n",
    "    def expected():\n",
    "        with R.dataflow():\n",
    "            C = R.ones([64, 128], \"int32\")\n",
    "\n",
    "            R.output(C)\n",
    "        return C\n",
    "\n",
    "    after = rewrite_call(pat, rewriter, before)\n",
    "    tvm.ir.assert_structural_equal(expected, after)\n",
    "\n",
    "\n",
    "def test_backtrack_for_no_op_rewriter_does_not_match_on_var():\n",
    "    \"\"\"The matches should always contain the bound value\n",
    "\n",
    "    This is a regression test.  In versions from\n",
    "    https://github.com/apache/tvm/pull/16732 to\n",
    "    https://github.com/apache/tvm/pull/16828, the `rewrite_call`\n",
    "    function could erroneously call the rewriter with `expr` and\n",
    "    `matches[pat]` set to a variable (`C`) instead of the value to\n",
    "    which it is bound (`R.add(A,B)`).\n",
    "    \"\"\"\n",
    "    pat_a = is_op(\"relax.add\")(wildcard(), wildcard())\n",
    "    pat_b = is_op(\"relax.add\")(wildcard(), wildcard())\n",
    "    pat = pat_a | pat_b\n",
    "\n",
    "    def rewriter(expr, matches):\n",
    "        assert isinstance(matches[pat], rx.Call)\n",
    "        return expr\n",
    "\n",
    "    @R.function(private=True)\n",
    "    def before():\n",
    "        with R.dataflow():\n",
    "            A = R.ones([64, 128], \"int32\")\n",
    "            B = R.zeros([64, 128], \"int32\")\n",
    "            C = R.add(A, B)\n",
    "\n",
    "            R.output(C)\n",
    "        return C\n",
    "\n",
    "    expected = before\n",
    "    after = rewrite_call(pat, rewriter, before)\n",
    "    tvm.ir.assert_structural_equal(expected, after)\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ai",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
